import datetime
import os
import typing as tp

from loguru import logger
from sqlalchemy import select, and_, func

from .signal import TaskEvent
from ..agent import Agent
from ..task import Task
from ..jobs.single_node_single_mod_job import Step as SNSMStep
from ..common import Result
from ..db.models.job import TriggeredTask, FinishedTask, DelayedTask, StepResult
from ..executors import BaseExecutor


@TaskEvent.start_signal.connect
async def task_start(sender: str, agent: Agent, task: Task):
    task.start_time = datetime.datetime.now()
    # 获取触发次数
    data = await agent.db.query(
        select(func.count(TriggeredTask.id))
        .where(and_(TriggeredTask.job_id == int(task.aps_job_id)))
    )
    index = data.scalar() + 1
    await agent.db.add(TriggeredTask(
        job_id=int(task.aps_job_id),
        trigger_index=index
    ))
    task.trigger_index = index

    task_log_dir_path = os.path.join(task.log_root_path, str(task.trigger_index))
    task_log_path = os.path.join(task_log_dir_path, "task.log")
    if not os.path.exists(task_log_dir_path):
        os.makedirs(task_log_dir_path)
    id_ = logger.add(
        task_log_path,
        level="DEBUG",
        filter=lambda r: r["message"].startswith(f"[{task.aps_job_id}"),  # important
        enqueue=True,
        # format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
        #        "<level>{level: <8}</level> | "
        #        "<level>{message}</level>",
    )
    task.loguru_log_path_id = id_
    task.log_path = task_log_path
    logger.info(f"[{task.aps_job_id}][{task.trigger_index}]{task.job} 已触发")
    if task.job.enable_viztracer:
        task.viztracer_result_json = os.path.join(task_log_dir_path, agent.cfg.base.viztracer_result_json)
        logger.info(f"[{task.aps_job_id}][{task.trigger_index}] 启动Viztracer -> {task.viztracer_result_json}")

@TaskEvent.finish_signal.connect
async def task_finish(sender: str, agent: Agent, task: Task):
    logger.info(f"[{task.aps_job_id}][{task.trigger_index}]{task.job} 已完成")
    logger.remove(task.loguru_log_path_id)
    await agent.db.add(FinishedTask(
        job_id=str(task.aps_job_id),
        trigger_index=task.trigger_index,
        start_time=task.start_time,
        end_time=datetime.datetime.now(),
        result=Result.NOT_IMPL.value.msg
    ))


@TaskEvent.executor_log_signal.connect
async def executor_log(sender: str, agent: Agent, task: Task, message: str, type_: tp.Literal["STDERR", "STDOUT"]):
    msg = f"[{task.aps_job_id}][{task.trigger_index}] | {message}"
    if type_ == "STDERR":
        logger.error(msg)
    else:
        logger.info(msg)


@TaskEvent.step_start_signal.connect
async def step_start(sender: str, agent: Agent, task: Task, step: tp.Union[SNSMStep]):
    task.current_step_start_time = datetime.datetime.now()
    logger.info(f"[{task.aps_job_id}][{task.trigger_index}] | [{task.trigger_index}、{step.name}] | {step.executor} | 已开始")


@TaskEvent.step_finish_signal.connect
async def step_finish(sender: str, agent: Agent, task: Task, step: tp.Union[SNSMStep], result: Result):
    await agent.db.add(StepResult(
        job_id=str(task.aps_job_id),
        trigger_index=task.trigger_index,
        step_index=task.current_step_index,
        start_time=task.current_step_start_time,
        end_time=datetime.datetime.now(),
        result=result.value.msg
    ))
    logger.info(f"[{task.aps_job_id}][{task.trigger_index}] | [{task.trigger_index}、{step.name}] | {step.executor} | 已完成，结果：{result.value}")
    task.current_step_index += 1


@TaskEvent.finally_signal.connect
async def finally_signal(sender: str, agent: Agent, task: Task):
    if task.is_delayed:
        await agent.db.add(DelayedTask(
            job_id=str(task.aps_job_id),
            trigger_index=task.trigger_index
        ))
